package core

import (
	"bytes"
	"fmt"
	"log"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
	"webssh/src/logger"

	"github.com/gorilla/websocket"
	"github.com/spf13/viper"
	"golang.org/x/crypto/ssh"
)

var (
	privatekeyPath = "/root/.ssh/id_rsa"
	publickeyPath  = "/root/.ssh/id_rsa.pub"
)

func NewSSHClient(tenantId string, username string) (SSHClient, error) {
	client := SSHClient{}
	client.Username = "root"
	client.Port = 22
	client.HomePath = userHomepath(tenantId, username)
	client.Username = username
	client.IPAddress = "localhost"
	return client, nil
}

func userHomepath(tenantId, username string) string {
	rootPath := viper.GetString("server.ceph-rootpath")
	userHomeSubpath := viper.GetString("server.ceph-userhome-subpath")
	return fmt.Sprintf("%s/%s/%s/%s", rootPath, tenantId, userHomeSubpath, username)
}

func readRootPublicKey() ([]byte, error) {
	b, err := os.ReadFile(publickeyPath)
	if err != nil {
		return nil, err
	}
	return b, nil
}

func addPublicKey(homeDir, publicKeyContent string) error {
	// 当前依赖freeipa管理用户，注入root用户公钥。
	sshDir := filepath.Join(homeDir, ".ssh")
	authorizedKeysPath := filepath.Join(sshDir, "authorized_keys")

	// 确保.ssh目录存在
	if err := os.MkdirAll(sshDir, 0700); err != nil {
		return err
	}

	// 读取已有的公钥
	var existingKeys []byte
	if b, err := os.ReadFile(authorizedKeysPath); err == nil {
		existingKeys = b
	} else if !os.IsNotExist(err) {
		return err
	}

	// 添加新公钥，如果已存在则不添加
	if !bytes.Contains(existingKeys, []byte(publicKeyContent)) {
		if err := os.WriteFile(authorizedKeysPath, append(existingKeys, []byte(publicKeyContent)...), 0600); err != nil {
			return err
		}
	}
	return nil
}

// GenerateClient 创建ssh客户端
func (sclient *SSHClient) GenerateClient() error {
	var (
		addr   string
		client *ssh.Client
		config ssh.Config
	)

	key, err := os.ReadFile(privatekeyPath)
	if err != nil {
		return fmt.Errorf("unable to read private key: %v", err)
	}
	signer, err := ssh.ParsePrivateKey([]byte(key))
	if err != nil {
		return fmt.Errorf("unable to parse private key: %v", err)
	}

	config = ssh.Config{
		Ciphers: []string{"aes128-ctr", "aes192-ctr", "aes256-ctr", "aes128-gcm@openssh.com", "arcfour256", "arcfour128", "aes128-cbc", "3des-cbc", "aes192-cbc", "aes256-cbc"},
	}
	clientConfig := &ssh.ClientConfig{
		User: sclient.Username,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		Timeout: 5 * time.Second,
		Config:  config,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}
	addr = fmt.Sprintf("%s:%d", sclient.IPAddress, sclient.Port)
	client, err = ssh.Dial("tcp", addr, clientConfig)
	if err != nil {
		// 判断是auth失败时向用户注入公钥重新连接
		logger.Info("tcp connect err, %v", err)
		if strings.Contains(err.Error(), "handshake failed") {
			// 向信任公钥列表中添加信任公钥
			rootPublickey, err := readRootPublicKey()
			if err != nil {
				return err
			}
			if err := addPublicKey(sclient.HomePath, string(rootPublickey)); err != nil {
				return err
			}
			client, err = ssh.Dial("tcp", addr, clientConfig)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}
	sclient.Client = client
	return nil
}

// InitTerminal 初始化终端
func (sclient *SSHClient) InitTerminal(ws *websocket.Conn, rows, cols int) *SSHClient {
	sshSession, err := sclient.Client.NewSession()
	if err != nil {
		log.Println(err)
		return nil
	}
	sclient.Session = sshSession
	sclient.StdinPipe, _ = sshSession.StdinPipe()
	wsOutput := new(wsOutput)
	sshSession.Stdout = wsOutput
	sshSession.Stderr = wsOutput
	wsOutput.ws = ws
	modes := ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

	if err := sshSession.RequestPty("xterm", rows, cols, modes); err != nil {
		return nil
	}
	if err := sshSession.Shell(); err != nil {
		return nil
	}
	return sclient
}

// Connect ws连接
func (sclient *SSHClient) Connect(ws *websocket.Conn, timeout time.Duration, closeTip string) {
	stopCh := make(chan struct{})
	go func() {
		// 配置客户端无操作，服务端session超时关闭
		timer := time.NewTimer(timeout)
		defer timer.Stop()
		defer ws.Close()
		for {
			select {
			case <-timer.C:
				log.Println("WebSocket connection timed out")
				return
			default:
				// p为用户输入
				_, p, err := ws.ReadMessage()
				if err != nil {
					close(stopCh)
					return
				}
				if string(p) == "ping" {
					ws.WriteMessage(websocket.TextMessage, nil)
					continue
				}

				if strings.Contains(string(p), "resize") {
					resizeSlice := strings.Split(string(p), ":")
					rows, _ := strconv.Atoi(resizeSlice[1])
					cols, _ := strconv.Atoi(resizeSlice[2])
					err := sclient.Session.WindowChange(rows, cols)
					if err != nil {
						log.Println(err)
						close(stopCh)
						return
					}
					continue
				}
				_, err = sclient.StdinPipe.Write(p)
				if err != nil {
					close(stopCh)
					return
				}
				timer.Reset(timeout)
			}
		}
	}()

	defer func() {
		ws.Close()
		sclient.Close()

		if err := recover(); err != nil {
			log.Println(err)
		}
	}()
	// 设置ws服务端超时时间timer
	stopTimer := time.NewTimer(timeout)
	defer stopTimer.Stop()
	// 主循环阻塞，超时后退出，defer中关闭websocket
	for {
		select {
		case <-stopCh:
			return
		case <-stopTimer.C:
			ws.WriteMessage(1, []byte(fmt.Sprintf("\u001B[33m%s\u001B[0m", closeTip)))
			return
		}
	}
}
